"""Prototype GPTL instrumentation helpers."""

from __future__ import annotations

import re
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional

# Add cpp-function-call-tree-mcp to path for access to the analyzer
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "cpp-function-call-tree-mcp" / "src"))

try:
    from cpp_analyzer import AnalysisConfig, CppCodeAnalyzer, FunctionInfo
    CPP_ANALYZER_AVAILABLE = True
except ImportError:
    CPP_ANALYZER_AVAILABLE = False

from .base import InstrumentationConfig, InstrumentationPatch, Instrumenter, should_process


@dataclass
class FunctionRegion:
    name: str
    start_line: int
    end_line: int
    indent: str
    kind: str  # "subroutine" or "function"
    instrument: bool


class FortranGPTLInstrumenter(Instrumenter):
    """Very early prototype for injecting GPTL calls into Fortran sources."""

    start_pattern = re.compile(r"^(?P<indent>\s*)(?P<kind>subroutine|function)\s+(?P<name>[a-z0-9_]+)", re.IGNORECASE)
    end_pattern = re.compile(r"^(?P<indent>\s*)end\s+(subroutine|function)(?:\s+(?P<name>[a-z0-9_]+))?", re.IGNORECASE)

    def plan(self, path: Path, config: InstrumentationConfig) -> InstrumentationPatch:
        original = path.read_text()
        modified = self._instrument_text(original, config)
        description = "Inserted GPTL start/stop pairs for Fortran routines"
        return InstrumentationPatch(path=path, original=original, modified=modified, description=description)

    def _instrument_text(self, source: str, config: InstrumentationConfig) -> str:
        lines = source.splitlines()
        stack: List[FunctionRegion] = []
        inserts: Dict[int, List[str]] = {}

        for idx, line in enumerate(lines):
            start_match = self.start_pattern.match(line)
            if start_match:
                name = start_match.group("name")
                kind = start_match.group("kind").lower()
                indent = start_match.group("indent") or ""
                region = FunctionRegion(
                    name=name,
                    start_line=idx,
                    end_line=-1,
                    indent=indent,
                    kind=kind,
                    instrument=should_process(name, config),
                )
                stack.append(region)
                continue

            end_match = self.end_pattern.match(line)
            if end_match and stack:
                region = stack.pop()
                region.end_line = idx
                end_name = end_match.group("name")
                if end_name and end_name.lower() != region.name.lower():
                    continue
                if region.instrument:
                    self._plan_fortran_inserts(region, lines, inserts)

        instrumented_lines: List[str] = []
        for idx, line in enumerate(lines):
            if idx in inserts:
                instrumented_lines.extend(inserts[idx])
            instrumented_lines.append(line)
        if len(lines) in inserts:
            instrumented_lines.extend(inserts[len(lines)])

        return "\n".join(instrumented_lines)

    def _plan_fortran_inserts(
        self,
        region: FunctionRegion,
        lines: List[str],
        inserts: Dict[int, List[str]],
    ) -> None:
        body_start = region.start_line + 1
        body_end = region.end_line

        if body_end <= body_start:
            return

        use_idx = self._locate_use_insert(body_start, body_end, lines)
        if not self._has_use_gptl(body_start, body_end, lines):
            indent = region.indent or ""
            inserts.setdefault(use_idx, []).append(f"{indent}  use GPTL")

        start_idx = self._locate_executable_insert(body_start, body_end, lines)
        indent = region.indent or ""
        inserts.setdefault(start_idx, []).append(f"{indent}  call GPTLstart('{region.name}')")

        stop_indent = region.indent or ""
        inserts.setdefault(region.end_line, []).append(f"{stop_indent}  call GPTLstop('{region.name}')")

    @staticmethod
    def _locate_use_insert(body_start: int, body_end: int, lines: List[str]) -> int:
        idx = body_start
        while idx < body_end:
            stripped = lines[idx].strip().lower()
            if not stripped or stripped.startswith("!"):
                idx += 1
                continue
            if stripped.startswith("use "):
                idx += 1
                continue
            break
        return idx

    @staticmethod
    def _locate_executable_insert(body_start: int, body_end: int, lines: List[str]) -> int:
        declaration_prefixes = (
            "use ",
            "implicit",
            "integer",
            "real",
            "double",
            "logical",
            "character",
            "type",
            "class",
            "complex",
            "dimension",
            "parameter",
            "procedure",
            "interface",
            "save",
            "data",
            "common",
        )
        idx = body_start
        while idx < body_end:
            stripped = lines[idx].strip().lower()
            if not stripped or stripped.startswith(("!", "#")):
                idx += 1
                continue
            if stripped.startswith(declaration_prefixes) or "::" in stripped.split("!")[0]:
                idx += 1
                continue
            break
        return idx

    @staticmethod
    def _has_use_gptl(body_start: int, body_end: int, lines: List[str]) -> bool:
        for idx in range(body_start, body_end):
            if re.search(r"^\s*use\s+gptl", lines[idx], re.IGNORECASE):
                return True
        return False


class CppGPTLInstrumenter(Instrumenter):
    """Enhanced GPTL injection for C/C++ translation units with AST support."""

    CONTROL_KEYWORDS = {"if", "for", "while", "switch", "catch"}
    function_pattern = re.compile(
        r"(?P<header>^[^{;]+?\b(?P<name>[A-Za-z_][A-Za-z0-9_]*)\s*\([^;{]*\)\s*)\{",
        re.MULTILINE,
    )

    # Enhanced pattern for single-line functions and member initializers
    single_line_function_pattern = re.compile(
        r"(?P<header>^[^{;]+?\b(?P<name>[A-Za-z_][A-Za-z0-9_]*)\s*\([^;{]*\)\s*)\{\s*(?P<body>[^{}]+)\s*\}",
        re.MULTILINE,
    )

    # Pattern for constructor member initializer lists
    constructor_initializer_pattern = re.compile(
        r"(?P<header>^[^{;]+?\b(?P<name>[A-Za-z_][A-Za-z0-9_]*)\s*\([^;{]*\)\s*):(?P<init_list>[^{]+)\{",
        re.MULTILINE,
    )

    def plan(self, path: Path, config: InstrumentationConfig) -> InstrumentationPatch:
        original = path.read_text()
        modified = self._instrument_text(original, config)
        description = "Inserted GPTL start/stop pairs for C/C++ functions"
        return InstrumentationPatch(path=path, original=original, modified=modified, description=description)

    def _instrument_text(self, source: str, config: InstrumentationConfig) -> str:
        # Try to use AST-based approach first
        if CPP_ANALYZER_AVAILABLE:
            try:
                return self._instrument_with_ast(source, config)
            except Exception:
                # Fall back to regex-based approach
                pass

        return self._instrument_with_regex(source, config)

    def _instrument_with_ast(self, source: str, config: InstrumentationConfig) -> str:
        """Use the MCP parser's AST for more accurate instrumentation."""
        # Create a temporary file for analysis
        import tempfile
        import os

        with tempfile.NamedTemporaryFile(mode='w', suffix='.cpp', delete=False) as tmp_file:
            tmp_file.write(source)
            tmp_path = tmp_file.name

        try:
            # Analyze the file using the MCP parser
            analysis_config = AnalysisConfig(
                paths=[tmp_path],
                include_source_code=True,
                use_clang=False  # Use regex parser which is more reliable
            )
            analyzer = CppCodeAnalyzer(analysis_config)
            result = analyzer.analyze()

            # Use the function information for precise instrumentation
            return self._instrument_with_function_info(source, result.functions, config)
        finally:
            os.unlink(tmp_path)

    def _instrument_with_function_info(self, source: str, functions: List, config: InstrumentationConfig) -> str:
        """Instrument using precise function information from AST."""
        instrumented = source
        offset = 0

        # Sort functions by start position for proper offset handling
        sorted_functions = sorted(
            [f for f in functions if hasattr(f, 'start_line') and f.start_line > 0],
            key=lambda f: f.start_line
        )

        lines = source.splitlines()

        for func in sorted_functions:
            if not should_process(func.name, config):
                continue

            if func.name in self.CONTROL_KEYWORDS:
                continue

            # Find function in source by line number
            if func.start_line <= len(lines) and func.end_line <= len(lines):
                func_start_idx = sum(len(lines[i]) + 1 for i in range(func.start_line - 1))
                func_end_idx = sum(len(lines[i]) + 1 for i in range(func.end_line)) - 1

                # Find opening brace
                brace_pos = source.find('{', func_start_idx)
                if brace_pos == -1 or brace_pos > func_end_idx:
                    continue

                # Insert start marker
                brace_pos += offset
                start_snippet = f"\n    GPTLstart(\"{func.name}\");"
                instrumented = instrumented[:brace_pos + 1] + start_snippet + instrumented[brace_pos + 1:]
                offset += len(start_snippet)

                # Find matching closing brace
                closing_brace = self._find_matching_brace(instrumented, brace_pos + 1 + len(start_snippet))
                if closing_brace is not None:
                    stop_snippet = f"    GPTLstop(\"{func.name}\");\n"
                    instrumented = instrumented[:closing_brace] + stop_snippet + instrumented[closing_brace:]
                    offset += len(stop_snippet)

        instrumented = self._ensure_include(instrumented)
        return instrumented

    def _instrument_with_regex(self, source: str, config: InstrumentationConfig) -> str:
        """Enhanced regex-based instrumentation supporting single-line functions."""
        instrumented = source
        offset = 0

        # First handle single-line functions
        single_line_matches = list(self.single_line_function_pattern.finditer(source))
        for match in single_line_matches:
            name = match.group("name")
            body = match.group("body").strip()

            if name in self.CONTROL_KEYWORDS:
                continue
            if not should_process(name, config):
                continue

            # Replace single-line function with multi-line instrumented version
            original_func = match.group(0)
            header = match.group("header")
            instrumented_func = f"{header}{{\n    GPTLstart(\"{name}\");\n    {body}\n    GPTLstop(\"{name}\");\n}}"

            start_pos = match.start() + offset
            end_pos = match.end() + offset
            instrumented = instrumented[:start_pos] + instrumented_func + instrumented[end_pos:]
            offset += len(instrumented_func) - len(original_func)

        # Handle constructor initializer lists
        init_matches = list(self.constructor_initializer_pattern.finditer(instrumented))
        for match in init_matches:
            name = match.group("name")
            if not should_process(name, config):
                continue

            # Find the opening brace after the initializer list
            brace_pos = match.end() + offset - 1  # Position of '{'
            start_snippet = f"\n    GPTLstart(\"{name}\");"
            instrumented = instrumented[:brace_pos + 1] + start_snippet + instrumented[brace_pos + 1:]
            offset += len(start_snippet)

            # Find matching closing brace
            closing_brace = self._find_matching_brace(instrumented, brace_pos + 1 + len(start_snippet))
            if closing_brace is not None:
                stop_snippet = f"    GPTLstop(\"{name}\");\n"
                instrumented = instrumented[:closing_brace] + stop_snippet + instrumented[closing_brace:]
                offset += len(stop_snippet)

        # Handle regular multi-line functions (avoiding already processed ones)
        # Remove single-line functions from consideration
        remaining_source = instrumented
        for match in reversed(single_line_matches):
            # Update positions based on offset
            start = match.start() + offset - len(source) + len(instrumented)
            end = match.end() + offset - len(source) + len(instrumented)
            if start >= 0 and end <= len(remaining_source):
                remaining_source = remaining_source[:start] + ' ' * (end - start) + remaining_source[end:]

        matches = list(self.function_pattern.finditer(remaining_source))
        for match in matches:
            name = match.group("name")
            if name in self.CONTROL_KEYWORDS:
                continue
            if not should_process(name, config):
                continue
            if not self._has_block_body(remaining_source, match.start(), match.end()):
                continue

            # Apply to the instrumented source
            body_start = match.end()
            start_snippet = f"\n    GPTLstart(\"{name}\");"
            instrumented = instrumented[:body_start] + start_snippet + instrumented[body_start:]
            offset += len(start_snippet)

            brace_index = self._find_matching_brace(instrumented, body_start + len(start_snippet))
            if brace_index is None:
                continue
            stop_snippet = f"    GPTLstop(\"{name}\");\n"
            instrumented = instrumented[:brace_index] + stop_snippet + instrumented[brace_index:]
            offset += len(stop_snippet)

        instrumented = self._ensure_include(instrumented)
        return instrumented

    @staticmethod
    def _has_block_body(source: str, start: int, end: int) -> bool:
        """Heuristic: only instrument when opening brace is followed by a newline."""

        brace_index = source.find("{", start, end + 1)
        if brace_index == -1:
            return False
        return brace_index + 1 < len(source) and source[brace_index + 1] in {"\n", "\r"}

    @staticmethod
    def _find_matching_brace(text: str, start_index: int) -> Optional[int]:
        depth = 1
        for idx in range(start_index, len(text)):
            ch = text[idx]
            if ch == '{':
                depth += 1
            elif ch == '}':
                depth -= 1
                if depth == 0:
                    return idx
        return None

    @staticmethod
    def _ensure_include(source: str) -> str:
        if "GPTLstart" not in source:
            return source
        if re.search(r"#\s*include\s*[<\"]gptl.h[>\"]", source):
            return source
        lines = source.splitlines()
        insert_idx = 0
        for idx, line in enumerate(lines):
            if line.startswith("#include"):
                insert_idx = idx + 1
        lines.insert(insert_idx, "#include \"gptl.h\"")
        return "\n".join(lines)


INSTRUMENTER_REGISTRY = {
    ("fortran", "gptl"): FortranGPTLInstrumenter(),
    ("cpp", "gptl"): CppGPTLInstrumenter(),
    ("c++", "gptl"): CppGPTLInstrumenter(),
}


def get_instrumenter(language: str, profile: str) -> Instrumenter:
    """Get instrumenter for language and profile, with platform support."""
    # First try platform-specific instrumenters
    try:
        from .platform_instrumenters import PLATFORM_INSTRUMENTER_REGISTRY
        key = (language.lower(), profile.lower())
        if key in PLATFORM_INSTRUMENTER_REGISTRY:
            return PLATFORM_INSTRUMENTER_REGISTRY[key]
    except ImportError:
        pass

    # Fall back to basic GPTL instrumenters
    key = (language.lower(), profile.lower())
    if key not in INSTRUMENTER_REGISTRY:
        raise ValueError(f"No instrumenter registered for language={language} profile={profile}")
    return INSTRUMENTER_REGISTRY[key]
